from text_embedding.chatlaw_text2vec import chatlaw_text2vec_encode
from text_embedding.m3e_base import m3e_encode
from text_embedding.bge_base import bge_encode
from text_embedding.text2vec_large import text2vec_encode
from text_embedding.multilingual_e5_small import e5_encode


embedding_name2dim_dict = {
    'chatlaw': 768,
    'm3e-base': 768,
    'bge-base': 768,
    'text2vec_large': 1024,
    'e5-small': 384
}

embedding_model_dict = {
    "chatlaw": chatlaw_text2vec_encode,
    "bge-base": bge_encode,
    "m3e-base": m3e_encode,
    "text2vec_large": text2vec_encode,
    "e5-small": e5_encode
}


def choose_model(embedding_type):
    return embedding_model_dict.get(embedding_type, m3e_encode)
